'''
    处理2列数据的绘图模块
        1.绘图函数后加_layout为相应的风格设置函数
'''
import plotly.graph_objs as go

# 两列都为字符串,绘制bubble图
def str_str_bubble(df,col1, col2):
    col1_labels=df[col1].unique()
    col2_labels=df[col2].unique()
    col2_len=len(col2_labels)
    x=[]
    #统计各点个数以决定bubble大小
    for lb in col1_labels:
        x+=[lb]*col2_len
    y=list(col2_labels)*len(col1_labels)
    x1,y1,size=[],[],[]
    #所有x,y相同位置的值对包括了所有可能的标签组合
    for a,b in zip(x,y):
        num=len(df[(df[col1]==a)&(df[col2]==b)])#该标签组合的数目
        if(num>0):
            x1.append(a)
            y1.append(b)
            size.append(num)
    trace = go.Scatter(
        x=x1,
        y=y1,
        mode='markers',
        marker=dict(
            size=size,
            sizemode='area',
            sizeref=2.*max(size)/(40.**2),
            sizemin=4
        )
    )
    return trace

def str_str_bubble_layout(col1,col2):
    title=','.join([col1,col2])
    layout=go.Layout(
                    title=title,
                    xaxis=dict(
                        title=col1
                    ),
                    yaxis=dict(
                        title=col2
                    ),
                )
    return layout

# 对两列数据绘制scatter
def two_column_scatter(df,col1,col2):
    x,y=df[col1],df[col2]
    trace = go.Scatter(
        x=x,
        y=y,
        mode='markers'
    )
    return trace

def two_column_scatter_layout(col1,col2):
    title=','.join([col1,col2])
    layout=go.Layout(
                    title=title,
                    xaxis=dict(
                        title=col1
                    ),
                    yaxis=dict(
                        title=col2
                    ),
                )
    return layout

# x轴为字符串,y轴为数值,绘制bubble
def str_num_bubble(df,str_col1, num_col2):
    col1_labels=df[str_col1].unique()
    col2_labels=df[num_col2].unique()
    col2_len=len(col2_labels)
    x=[]
    #统计各点个数以决定bubble大小
    for lb in col1_labels:
        x+=[lb]*col2_len
    y=list(col2_labels)*len(col1_labels)
    x1,y1,size=[],[],[]
    #所有x,y相同位置的值对包括了所有可能的标签组合
    for a,b in zip(x,y):
        num=len(df[(df[str_col1]==a)&(df[num_col2]==b)])#该标签组合的数目
        if(num>0):
            x1.append(a)
            y1.append(b)
            size.append(num)
    trace = go.Scatter(
        x=x1,
        y=y1,
        mode='markers',
        marker=dict(
            size=size,
            sizemode='area',
            sizeref=2.*max(size)/(40.**2),
            sizemin=4
        )
    )
    return trace

def str_num_bubble_layout(col1,col2):
    title=','.join([col1,col2])
    layout=go.Layout(
                    title=title,
                    xaxis=dict(
                        title=col1
                    ),
                    yaxis=dict(
                        title=col2
                    ),
                )
    return layout
